#include "stdafx.h"
#include "bvh.h"

#define MAX_DEPTH 100
//sah true or false
#define SAH 0

namespace bvh_build_internal
{
	struct BuildStateStruct
	{
		size_t segmentStart, segmentEnd;
		size_t nodeIndex;
		BBox centroidBBox;

		int depth;
	};

	struct CentroidWithID
	{
		Point centroid;
		size_t origIndex;
	};
}

using namespace bvh_build_internal;

//An iterative split in the middle build for BVHs
void BVH::build(const std::vector<Primitive*> &_objects)
{
	//all bounding boxes of all objects
	std::vector<BBox> objectBBoxes(_objects.size());

	BuildStateStruct curState;
	curState.centroidBBox = BBox::empty();

	curState.depth = 0;

	//centroids of all primitives in the scene
	std::vector<CentroidWithID> centroids(objectBBoxes.size());

	//for SAH-calculation
	scene = BBox::empty();
	
	//initialize the bounding box vector
	for(size_t i = 0; i < objectBBoxes.size(); i++)
	{
		objectBBoxes[i] = _objects[i]->getBBox();
		
	    centroids[i].centroid = objectBBoxes[i].min.lerp(objectBBoxes[i].max, 0.5f);
	    centroids[i].origIndex = i;
	    curState.centroidBBox.extend(centroids[i].centroid);
		
		//for SAH-calculation
		scene.extend(objectBBoxes[i]);
	}

	//initialize counters for pseudo recursion
	curState.segmentStart = 0;
	curState.segmentEnd = objectBBoxes.size();
	curState.nodeIndex = 0;
	//prepare the bvh-tree
	m_nodes.resize(1);

	//build the stack for the pseudo recursion
	std::stack<BuildStateStruct> buildStack;

	const float _EPS = 0.0000001f;
	
	//flag if sah founds a splitvalue, if not make child
	bool foundcost = true;

	int maxdepth = 0;
	

	//endless loop, which will be stopped by 'break' when stack gets empty
	for(;;)
	{
		//number of the objects in the actual part
	    size_t objCnt = curState.segmentEnd - curState.segmentStart;
		Vector boxDiag = curState.centroidBBox.diagonal();
		
		//split the actual box at the longest side
		int splitDim = boxDiag.x > boxDiag.y ? (boxDiag.x > boxDiag.z ? 0 : 2) : (boxDiag.y > boxDiag.z ? 1 : 2);

		if(fabs(boxDiag[splitDim]) < _EPS || objCnt < 3 || curState.depth>=MAX_DEPTH || !foundcost)
		{
			foundcost = true;
			maxdepth = std::max(maxdepth, curState.depth);
			//Create a leaf
			const size_t NODE_TYPE_MASK = ((size_t)1 << Node::LEAF_FLAG_BIT);
			m_nodes[curState.nodeIndex].bbox = BBox::empty();
			m_nodes[curState.nodeIndex].dataIndex = m_leafData.size() | NODE_TYPE_MASK;

			//save all objects of the actual segment in this leafnode
			for(size_t i = curState.segmentStart; i < curState.segmentEnd; i++)
			{
				m_nodes[curState.nodeIndex].bbox.extend(objectBBoxes[centroids[i].origIndex]);
				m_leafData.push_back(_objects[centroids[i].origIndex]);
			}

			m_leafData.push_back(NULL);

			//if all objects are saved in the tree
			if(buildStack.empty())
				break;

			//else take the next unfinished part from the stack
			curState = buildStack.top();
			buildStack.pop();

			continue;
		}

		//calculate and save the position of the split-plane
		//without SAH simply split in the middle of the longest edge
		float splitVal = (curState.centroidBBox.min[splitDim] + curState.centroidBBox.max[splitDim]) / 2.f;

		
		std::stack<size_t> bestleftStack;
		std::stack<size_t> bestrightStack;

		//SAH calculation
		if(SAH==1){
			foundcost=false;

			//calculate sah split-plane
			float height[3] = {curState.centroidBBox.max[2]-curState.centroidBBox.min[2],curState.centroidBBox.max[0]-curState.centroidBBox.min[0],curState.centroidBBox.max[1]-curState.centroidBBox.min[1]};
			float depth[3] = {curState.centroidBBox.max[1]-curState.centroidBBox.min[1],curState.centroidBBox.max[2]-curState.centroidBBox.min[2],curState.centroidBBox.max[0]-curState.centroidBBox.min[0]};
			float bestcost = FLT_MAX;
			float bestpos = 0;
			int bestDim = 0,
				bestPrim = 0;

			//work on a copy of the actual State stack
			float lastWidth[3]={0,0,0};
			
			float width[3], width2[3], split[3], gap[3];
			for(int i=0;i<3;i++){
				gap[i]= (curState.centroidBBox.max[i]-curState.centroidBBox.min[i])/100.f;
				width[i]= curState.centroidBBox.min[i]-gap[i];
			
			}

			//test the splitVal for every raster position
			while((curState.centroidBBox.min[1]+width[1])< curState.centroidBBox.max[1]){
				for(int i=0;i<3;i++){
					width[i]+= gap[i];
				}
				//compute the width of both parts
				for(int i=0;i<3;i++)
				{
					width2[i] = curState.centroidBBox.max[i]-curState.centroidBBox.min[i]-width[i];
					split[i] = curState.centroidBBox.min[i] + width[i];
				}
				//if(width[0]<=0 || width[1]<=0 || width[2]<=0 || width2[0]<=0 || width2[1]<=0 || width2[2]<=0)continue; //the leftest primitive can be smaller than the state.min and it is not possible/necessary to split at this position.
				//if the distance to previous computed splitVal is too small, skip this one
				/*if(bestcost!=FLT_MAX
					&& (fabs(split[0]-lastWidth[0])<(curState.centroidBBox.diagonal()[0]/100.f)
						|| fabs(split[1]-lastWidth[1])<(curState.centroidBBox.diagonal()[1]/100.f)
						|| fabs(split[2]-lastWidth[2])<(curState.centroidBBox.diagonal()[2]/100.f)) )
						continue;*/
				
				//calculate the areas under the parts
				float area[3], area2[3], totarea[3];
				for(int i=0;i<3;i++){
					area[i] = 2*(width[i]*height[i]+depth[i]*height[i]+width[i]*depth[i]);
					area2[i] = 2*(width2[i]*height[i]+depth[i]*height[i]+width2[i]*depth[i]);
					totarea[i]= 2*((width2[i]+width[i])*height[i]+depth[i]*height[i]+(width2[i]+width[i])*depth[i]);
				}
				
				//create a copy of the actual State stack to extract how many objects are on the right/left side
				std::stack<size_t> leftStack[3];
				std::stack<size_t> rightStack[3];

				for(size_t l=curState.segmentStart;l<curState.segmentEnd;l++){
					size_t curcountPrim = l;

					//left-right-split Dim 0..2
					for(int i=0;i<3;i++){
						//if primitive is on the left side of the split plane
						if(centroids[curcountPrim].centroid[i] <= split[i])
						{
							leftStack[i].push(curcountPrim);
						}
						else 
							//if primitive is on the right side of the split plane
							{
								rightStack[i].push(curcountPrim);
							}
					}
				}

				//compute the number of elements in the left- and rightstacks
				size_t leftcount[3], rightcount[3];
				for(int i=0;i<3;i++){
					leftcount[i] = leftStack[i].size();
					rightcount[i] = rightStack[i].size();
				}
					
				//the real costs of this splitVal
				float cost[3];
				for(int i=0;i<3;i++){
					if(leftcount[i]!=0 && rightcount[i]!=0)
						cost[i] = (float)(((area[i]/totarea[i])*(leftcount[i]))+((area2[i]/totarea[i])*(rightcount[i])));
					else
						cost[i]=FLT_MAX;
				}

				//if in one part are zero faces, just skip this splitplane
				
				//if the actual costs are better, save them
				for(int i=0;i<3;i++){
					if( (curState.segmentEnd-curState.segmentStart)!=leftcount[i] && (curState.segmentEnd-curState.segmentStart)!=rightcount[i]				//child-nodes contain other #objects as actual node
											//if object-size is small, you can split wherever you want
																	//if objects-size is bigger, the child-nodes must have a min-ratio
						)
					if(width[i]>0 && width2[i]>0 )
						//if cost are better then before, save all the properties of this split
						if(cost[i]<bestcost){
							bestcost = cost[i];
							bestpos = split[i];
							bestDim = i;
							bestleftStack = leftStack[i];
							bestrightStack = rightStack[i];
							//std::cout<<"area: "<<area[bestDim]<<"-"<<bestleftStack.size()<<"-"<<bestrightStack.size()<<"-"<<area2[bestDim]<<"-"<<bestcost<<std::endl;
							foundcost = true;
							//std::cout<<"area: "<<area[i]<<" lsize: "<<leftStack[i].size()<<"area2: "<<area2[i]<<" rsize: "<<rightStack[i].size()<<std::endl;
						}
					}
				//save the actual splitPositions, to avoid calculation of splitplanes with too little space between
				lastWidth[0] = split[0];
				lastWidth[1] = split[1];
				lastWidth[2] = split[2];
				}

			if(!foundcost) continue;
			splitDim=bestDim;
			splitVal=bestpos;
			
			//std::cout<<"whole: "<<curState.segmentEnd-curState.segmentStart<<" left: "<<bestleftStack.size()<<" right: "<<bestrightStack.size()<<" splitDim: "<<bestDim<<std::endl;
		}
		//end of SAH-calculation

		size_t leftPtr = curState.segmentStart,
			rightPtr = curState.segmentEnd - 1;


		BBox nodeBBox = BBox::empty();
		BuildStateStruct rightState;
		curState.centroidBBox = BBox::empty();
		rightState.centroidBBox = BBox::empty();

		//handle all objects in the actual node
		//split the objects into the ones on the left side and on the right side
		//and create the BBoxes for these nodes
		while(leftPtr < rightPtr)
		{
			while(leftPtr < curState.segmentEnd && centroids[leftPtr].centroid[splitDim] <= splitVal)
			{
				//extend the left part with objects
				curState.centroidBBox.extend(centroids[leftPtr].centroid);
				nodeBBox.extend(objectBBoxes[centroids[leftPtr].origIndex]);
				leftPtr++;
			}

			while(rightPtr >= curState.segmentStart && centroids[rightPtr].centroid[splitDim] > splitVal)
			{
				//extend the right part with objects
				rightState.centroidBBox.extend(centroids[rightPtr].centroid);
				nodeBBox.extend(objectBBoxes[centroids[rightPtr].origIndex]);
				rightPtr--;
			}

			//if all objects were handled
			if(leftPtr >= rightPtr)
				break;

			std::swap(centroids[leftPtr], centroids[rightPtr]);
		}

		m_nodes[curState.nodeIndex].bbox = nodeBBox;
		m_nodes[curState.nodeIndex].dataIndex = m_nodes.size();

		_ASSERT(leftPtr > curState.segmentStart && leftPtr < curState.segmentEnd);
		_ASSERT(leftPtr == rightPtr + 1);

		//set the pointers for the new children
		rightState.segmentStart = leftPtr;
		rightState.segmentEnd = curState.segmentEnd;
		curState.segmentEnd = leftPtr;
		curState.nodeIndex = m_nodes.size();
		rightState.nodeIndex = curState.nodeIndex + 1;

		//increase depth for the children
		curState.depth = curState.depth+1;
		rightState.depth = curState.depth;
		
		buildStack.push(rightState);

		m_nodes.resize(rightState.nodeIndex + 1);
	}
	std::cout<<"max Depth: "<<maxdepth<<std::endl;
	if(SAH==1)
		std::cout<<"Build with BVH with SAH "<<std::endl;
	else
		std::cout<<"Build with BVH without SAH "<<std::endl;
}

//Recursive intersection
BVH::IntersectionReturn BVH::intersect(const Ray &_ray, float _previousBestDistance) const
{
	Primitive::IntRet bestHit;
	bestHit.distance = _previousBestDistance;

	Primitive *bestPrimitive = NULL;

	//stack with the elements which must get handled
	std::stack<size_t> traverseStack;

	size_t curNode = 0;
	//size_t closestFace = (size_t)-1;		this variable isn't used anymore

	for(;;)
	{
		const BVH::Node& node = m_nodes[curNode];
		//if the actual node is a leaf
		if(node.isLeaf())
		{
			//take the first child of the leaf
			size_t idx = node.getLeftChildOrLeaf();
			//while loop over all children of the actual node
			while(m_leafData[idx] != NULL)
			{
				//intersect the actual child with the ray
				Primitive::IntRet curRet = m_leafData[idx]->intersect(_ray, bestHit.distance);
				//if an intersection occurs and the distance to this intersection is shorter than the actual best
				if(curRet.distance > Primitive::INTEPS() && curRet.distance < bestHit.distance)
				{
					//save the better distance and hitten object
					bestHit = curRet;
					bestPrimitive = m_leafData[idx];
				}
				//increase counter for the childindex
				idx++;
			}
			//if all parts of the tree were handled
			if(traverseStack.empty())
				break;
			//else take the next part from the stack
			curNode = traverseStack.top();
			traverseStack.pop();
		}
		//if the actual part is an inner node
		else
		{
			//get both children
			const BVH::Node &leftNode = m_nodes[node.getLeftChildOrLeaf()];
			const BVH::Node &rightNode = m_nodes[node.getLeftChildOrLeaf() + 1];

			//intersect with each BBox of them
			std::pair<float, float> intLeft = leftNode.bbox.intersect(_ray);
			std::pair<float, float> intRight = rightNode.bbox.intersect(_ray);
			intLeft.first = std::max(Primitive::INTEPS(), intLeft.first);
			intRight.first = std::max(Primitive::INTEPS(), intRight.first);
			intLeft.second = std::min(intLeft.second, bestHit.distance);
			intRight.second = std::min(intRight.second, bestHit.distance);

			//booleans, whether the BBoxes were hitten or not, or are too close
			bool descendLeft = intLeft.first < intLeft.second + Primitive::INTEPS();
			bool descendRight = intRight.first < intRight.second + Primitive::INTEPS();

			//if only the left child is hitten, continue with this one
			if(descendLeft && !descendRight)
				curNode = node.getLeftChildOrLeaf();
			//if only the right child is hitten, continue with this one
			else if(descendRight && !descendLeft)
				curNode = node.getLeftChildOrLeaf() + 1;
			//if both children were hitten
			else if(descendLeft && descendRight)
			{
				//now the actual node is the left child of the node
				curNode = node.getLeftChildOrLeaf();
				size_t farNode = curNode + 1;
				//if right child is hitten before the left child, change the actual node to the right child of the node
				if(intLeft.first > intRight.first)
					std::swap(curNode, farNode);
				//put the farer away node on the stack
				traverseStack.push(farNode);
			}
			//if no one of the children is hitten
			else
			{
				//test if the stack is empty
				if(traverseStack.empty())
					break;
				//skip handling of this part of the tree and continue with the next part on the stack
				curNode = traverseStack.top();
				traverseStack.pop();
			}
		}
	}
	//create the return struct and save the distance and the hitten primitive in it
	BVH::IntersectionReturn ret;
	ret.ret = bestHit;
	ret.primitive = bestPrimitive;
	return ret;
}

